{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "08517b2e",
   "metadata": {},
   "source": [
    "# Result-Dependent if/else\n",
    "\n",
    "To branch based on the result from an electron, put the branch (if/else) logic inside another electron.\n",
    "\n",
    "## Context\n",
    "\n",
    "Often the output of one task is a value used to choose the execution path. In these cases, perform the if/else logic inside an electron. \n",
    "\n",
    "When a lattice is dispatched, the Covalent server executes the lattice in order to build the transport graph. The transport graph is then analyzed to parallelize electron execution on their assigned executors.\n",
    "\n",
    "If the server encounters a branch decision based on the output of an electron, it cannot infer the structure on which the decision depends and is prevented from building the transport graph.\n",
    "\n",
    "## Best Practice\n",
    "\n",
    "Compute branching inside an electron. Electrons' execution is deferred during the graph build phase, so their output cannot be used to build the transport graph and analyze the execution for parallelization. Instead, the electron is added to the transport graph and the branching is computed within the electron when it is executed.\n",
    "\n",
    "For result-dependent computations that might be too complex to encapsulate in a single electron, use a [sublattice](./dynamic_workflow.ipynb).\n",
    "\n",
    "## Example\n",
    "\n",
    "Contrast the two examples below.\n",
    "\n",
    "### Example 1: Incorrect\n",
    "\n",
    "This example demonstrates the incorrect approach: choosing the execution path in the lattice based on the output of `task_1`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "b8276e3f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Lattice Result\n",
      "==============\n",
      "status: POSTPROCESSING_FAILED\n",
      "result: None\n",
      "input args: ['5']\n",
      "input kwargs: {}\n",
      "error: Post-processing failed: Traceback (most recent call last):\n",
      "  File \"/Users/dave/agnostiq/covalent/covalent/executor/utils/wrappers.py\", line 36, in io_wrapper\n",
      "    output = fn(*args, **kwargs)\n",
      "  File \"/Users/dave/agnostiq/covalent/covalent/executor/base.py\", line 92, in wrapper_fn\n",
      "    output = fn(*new_args, **new_kwargs)\n",
      "  File \"/Users/dave/agnostiq/covalent/covalent_dispatcher/_core/runner.py\", line 433, in _post_process\n",
      "    result = workflow_function(*args, **kwargs)\n",
      "  File \"/var/folders/l_/bv3pdx7142df5_sht34x64q80000gn/T/ipykernel_10514/3352216439.py\", line 19, in workflow\n",
      "  File \"/Users/dave/agnostiq/covalent/covalent/_workflow/electron.py\", line 640, in wrapper\n",
      "    return electron_object(*args, **kwargs)\n",
      "  File \"/Users/dave/agnostiq/covalent/covalent/_workflow/electron.py\", line 321, in __call__\n",
      "    id, output = active_lattice.electron_outputs[0]\n",
      "IndexError: list index out of range\n",
      "\n",
      "\n",
      "start_time: 2023-03-13 21:14:56.656984\n",
      "end_time: 2023-03-13 21:14:56.785618\n",
      "\n",
      "results_dir: /Users/dave/.local/share/covalent/data\n",
      "dispatch_id: 6bcc035a-08e4-4077-a5a6-2076d7c3841c\n",
      "\n",
      "Node Outputs\n",
      "------------\n",
      "task_1(0): 10\n",
      ":parameter:5(1): 5\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import covalent as ct\n",
    "\n",
    "# Technique 1: Incorrect\n",
    "\n",
    "@ct.electron\n",
    "def task_1(x):\n",
    "    return x * 2\n",
    "\n",
    "@ct.electron\n",
    "def task_2(x):\n",
    "    return x ** 3\n",
    "\n",
    "@ct.lattice\n",
    "def workflow(a):\n",
    "    \n",
    "    res_1 = task_1(a)\n",
    "\n",
    "    if res_1 == 10:\n",
    "        final_res = task_2(res_1)\n",
    "    else:\n",
    "        final_res = res_1\n",
    "    \n",
    "    return final_res\n",
    "\n",
    "id = ct.dispatch(workflow)(5)\n",
    "res = ct.get_result(id, wait=True)\n",
    "print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce7ff657",
   "metadata": {},
   "source": [
    "### Example 2: Improved\n",
    "\n",
    "The output of `task_1` is passed to the `task_2_new`, which executes the chosen path internally and returns the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "7cc24313",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Lattice Result\n",
      "==============\n",
      "status: COMPLETED\n",
      "result: 1000\n",
      "input args: ['5']\n",
      "input kwargs: {}\n",
      "error: None\n",
      "\n",
      "start_time: 2023-03-13 21:14:58.815753\n",
      "end_time: 2023-03-13 21:14:58.977721\n",
      "\n",
      "results_dir: /Users/dave/.local/share/covalent/data\n",
      "dispatch_id: c16ebb6f-1b3b-43c7-a95f-fd16334c8dc9\n",
      "\n",
      "Node Outputs\n",
      "------------\n",
      "task_1(0): 10\n",
      ":parameter:5(1): 5\n",
      "task_2_new(2): 1000\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import covalent as ct\n",
    "\n",
    "# Technique 2: Correct\n",
    "\n",
    "@ct.electron\n",
    "def task_1(x):\n",
    "    return x * 2\n",
    "\n",
    "# Method (2):\n",
    "@ct.electron\n",
    "def task_2_new(x):\n",
    "    if x == 10:\n",
    "        return x ** 3\n",
    "    else:\n",
    "        return x\n",
    "\n",
    "@ct.lattice\n",
    "def workflow_2(a):\n",
    "    res_1 = task_1(a)\n",
    "    return task_2_new(res_1)\n",
    "\n",
    "id = ct.dispatch(workflow_2)(5)\n",
    "res = ct.get_result(id, wait=True)\n",
    "print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eee2041a",
   "metadata": {},
   "source": [
    "## See Also\n",
    "\n",
    "[Result-Dependent Loops](./result_dependent_loop.ipynb)\n",
    "\n",
    "[Dynamic Workflows](./dynamic_workflow.ipynb)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
